import numpy as np
import math
import os
import redis
import umsgpack
import pandas as pd

from settings import REDIS_CON, PRE_PATH

embed_data = []


def get_embed_data():
    if len(embed_data) == 0:
        redis_conn = redis.Redis(**REDIS_CON)
        # dataframe转化为hash键值对格式，获取存量数据已经序列化保存到redis内存的数据
        retrieved_data = redis_conn.hget('total_embed_tag', 'umsgpack_data')
        # 反序列化为字典对象
        deserialized_df_dict = umsgpack.unpackb(retrieved_data)
        # 需要将字典转换为DataFrame
        df = pd.DataFrame(deserialized_df_dict)
        embed_data.append(df)
    return embed_data[0]


class ImageEmbedProduce(object):
    def calculate_image_num(self, video_sec):
        # 根据视频时长计算视频需要的图片数量，向上取整
        total_num = math.ceil(video_sec // 3)
        if total_num < 10:
            total_num = 10
        expand_num = total_num * 3
        return total_num, expand_num

    def cosine_similarity(self, a, b):
        # 定义余弦相似度函数，进行计算
        result = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
        return result

    # search through the reviews for a specific product
    def top_n_similarity(self, df, user_query, top_n=8):
        df["similarities"] = df.embeddings_labels.apply(lambda x: self.cosine_similarity(x, user_query))
        result_df = df.sort_values("similarities", ascending=False)
        # 判断图片的md5是否重复，如果重复保留排名靠前的图片
        result_df = result_df.drop_duplicates(subset=['image_md5'])
        top_df = result_df.head(top_n)
        return top_df

    def get_image_path(self, result_df):
        image_paths = [os.path.join(PRE_PATH, path) for path in result_df['local_path'].values]
        return image_paths

    def run(self, video_sec, embedding):
        embed_df = get_embed_data()
        image_num, expand_num = self.calculate_image_num(video_sec)
        # 字幕向量转化为浮点数
        input_embedding_float = [float(value) for value in embedding]
        res_df = self.top_n_similarity(embed_df, input_embedding_float, top_n=expand_num)
        expand_image_paths = self.get_image_path(res_df)
        # 从扩展的图片中随机抽取需要的图片
        image_paths = np.random.choice(expand_image_paths, image_num, replace=False)
        return image_paths


if __name__ == '__main__':
    video_sec = 42.266666666666666
    # 简单尝试是否自定义余弦进行计算
    input_embedding = [0.1, 0.2, 0.3, 0.4, 0.5]
    input_embedding_float = [float(value) for value in input_embedding]

